package com.github.cyf.config;

import com.github.cyf.common.util.MDCUtil;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.concurrent.*;

/**
 * 方式一：自定义 Mdc 线程池
 * @author chenyifan
 * @create 2024-09-02 18:11
 */
public class MdcThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {

    public MdcThreadPoolTaskExecutor() {
    }

    @Override
    public void execute(Runnable command) {
        super.execute(wrap(command, MDCUtil.getCopyOfContextMap()));
    }

    @Override
    public Future<?> submit(Runnable task) {
        return super.submit(wrap(task, MDCUtil.getCopyOfContextMap()));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return super.submit(wrap(task, MDCUtil.getCopyOfContextMap()));
    }

    private static <T> Callable wrap(Callable<T> task, Map<String, String> contextMap) {
        return () -> {
            try {
                MDCUtil.setContextMap(contextMap);
                return task.call();
            } catch (Exception e) {
                throw new RuntimeException(e);
            } finally {
                MDCUtil.clear();
            }
        };
    }

    private static Runnable wrap(Runnable task, Map<String, String> contextMap) {
        return () -> {
            try {
                MDCUtil.setContextMap(contextMap);
                task.run();
            } finally {
                MDCUtil.clear();
            }
        };
    }
}
